﻿#pragma once

#include <functional>

#include <xy/math/math.h>
#include <xy/math/matrix.h>

namespace xy
{

namespace math
{

struct diff
{
    /**
     * @brief 获得导函数
     *
     * @tparam T
     * @param[in] f
     * @return std::function<T(double)>
     */
    template <class T> static std::function<T(double)> differential(std::function<T(double)> f)
    {
        static constexpr double h = numeric::distance;
        return [=](double x) { return (f(x + h) - f(x - h)) / (2 * h); };
    }

    /**
     * @brief 多元函数的 Jacobian 矩阵
     *
     * @tparam T
     * @param[in] f
     * @param[in] x
     * @return matxd
     */
    template <class T> static matxd jacobian(std::function<T(const vecxd &)> f, const vecxd &x)
    {
        // 巨坑：有些版本编译器不支持将 static_assert(false) 写在 else 里
        static constexpr double h = numeric::distance;
        static_assert(std::is_same_v<T, vecxd> || std::is_same_v<T, double>, "Unsupported type.");

        T y = f(x);
        std::size_t m = 0;
        if constexpr (std::is_same_v<T, vecxd>)
            m = y.size();
        else if constexpr (std::is_same_v<T, double>)
            m = 1;

        std::size_t n = x.size();
        matxd d(m, n);
        for (std::size_t i = 0; i < n; i++)
        {
            vecxd y1 = x, y2 = x;
            y1[i] += h, y2[i] -= h;

            // 中心差分向量
            T dp = (f(y1) - f(y2)) / (2 * h);
            if constexpr (std::is_same_v<T, vecxd>)
            {
                for (std::size_t j = 0; j < m; j++)
                    d(j, i) = dp[j];
            }
            else if constexpr (std::is_same_v<T, double>)
                d(0, i) = dp;
        }
        return d;
    }

    /**
     * @brief 多元函数的梯度向量
     *
     * @tparam T
     * @param[in] f
     * @param[in] x
     * @return matxd
     */
    template <class T> static matxd gradient(std::function<T(const vecxd &)> f, const vecxd &x)
    {
        return jacobian(f, x).transpose();
    }

    /**
     * @brief 多元函数的 Hessian 矩阵
     *
     * @param[in] f
     * @param[in] x
     * @return matxd
     */
    static matxd hessian(std::function<double(const vecxd &)> f, const vecxd &x)
    {
        static constexpr double h = numeric::distance;

        auto d = jacobian(f, x);
        std::size_t n = x.size();
        matxd H(n, n);
        for (std::size_t i = 0; i < n; i++)
        {
            vecxd y1 = x, y2 = x;
            y1[i] += h, y2[i] -= h;

            auto dp = (jacobian(f, y1) - jacobian(f, y2)) / (2 * h);
            for (std::size_t j = 0; j < n; j++)
                H(i, j) = dp(0, j);
        }
        return H;
    }
};

} // namespace math

} // namespace xy
